
package com.ctrip.corp.obt.shard.sql.mybatis.interceptor.inner;

import static com.ctrip.corp.obt.generic.utils.Conditional.ofNullable;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.sql.Connection;
import java.text.MessageFormat;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Properties;
import java.util.Set;
import java.util.TreeSet;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.session.Configuration;

import com.ctrip.corp.obt.generic.core.context.RequestContext;
import com.ctrip.corp.obt.generic.utils.CollectionUtils;
import com.ctrip.corp.obt.generic.utils.JsonUtils;
import com.ctrip.corp.obt.metric.Metrics;
import com.ctrip.corp.obt.metric.log.model.MetricMeta;
import com.ctrip.corp.obt.metric.spectator.api.Id;
import com.ctrip.corp.obt.shard.sql.parser.AbstractSqlParserSupport;
import com.ctrip.corp.obt.shard.sql.parser.IPluginPropertyParserRequest;
import com.ctrip.corp.obt.shard.sql.parser.forbidden.UpdateForbiddenMatcher;
import com.ctrip.corp.obt.shard.sql.parser.forbidden.UpdateForbiddenMatcherComposite;
import com.ctrip.corp.obt.shard.utils.ClassUtils;
import com.ctrip.corp.obt.shard.utils.ParameterPropertyUtils;
import com.ctrip.corp.obt.shard.utils.PluginUtils;
import com.ctrip.corp.obt.shard.utils.PluginUtils.DSBoundSql;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.JdbcParameter;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.expression.operators.relational.ItemsListVisitor;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.expression.operators.relational.NamedExpressionList;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.SubSelect;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.statement.update.UpdateSet;

/**
 * 对update语句中set字段进行处理
 *
 * @author xukang zhang
 * @date 2023-08-16
 */
@Slf4j
@Data
@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
public class UpdateForbiddenProcessPlugin extends AbstractSqlParserSupport implements InnerPlugin {

    private static final String UNKNOWN = "UNKNOWN";
    private static final String TABLE = "table";
    private static final String REQUEST_URI = "requestUri";
    private static final String SECURITY_DAL_FORBIDDEN = "arch.shard.forbidden";
    private static final String CRLF = "\r\n";
    private static final Id SHARD_FORBIDDEN_SINGLE = Metrics.REGISTRY.createId(SECURITY_DAL_FORBIDDEN + ".single");
    private static final Id SHARD_FORBIDDEN_BATCH = Metrics.REGISTRY.createId(SECURITY_DAL_FORBIDDEN + ".batch");
    private static final Id SHARD_FORBIDDEN_INSERT = Metrics.REGISTRY.createId(SECURITY_DAL_FORBIDDEN + ".insert");
    private static final String LOGGER_SUFFIX = ".logger";
    private static final String LIST_PREFIX_NAME = "__frch_item_";
    private static final Pattern EXTRACT_LIST_INFO_PATTERN = Pattern.compile("^__frch_item_(\\d+)\\.(.+)$");
    private static final List<String> LIST_ALIAS = Arrays.asList("list", "collection", "arg0");
    private static final String FORBIDDEN_MATCHER = "forbiddenMatcher";
    private static final String REMOVED_FLAG = "removed";
    private static final String SUCCESS_FLAG = "success";
    private UpdateForbiddenMatcherComposite updateForbiddenMatcherComposite = new UpdateForbiddenMatcherComposite();

    /**
     * update sql process
     */
    @Override
    protected void processUpdate(Update update, int index, String sql, Object obj) {
        // validate the parameters
        if (!(obj instanceof MultiParserRequest)) {
            return;
        }
        Instant currentStartTime = Instant.now();
        MultiParserRequest multiParserRequest = (MultiParserRequest)obj;
        if (!IPluginPropertyParserRequest.validate(multiParserRequest)) {
            return;
        }
        PluginUtils.DSBoundSql boundSql = multiParserRequest.getBoundSql();
        Object parameterObject = boundSql.parameterObject();
        List<ParameterMapping> parameterMappings = boundSql.parameterMappings();
        ArrayList<UpdateSet> updateSets = update.getUpdateSets();
        if (CollectionUtils.isEmpty(updateSets) || parameterObject == null) {
            return;
        }
        Set<ParameterMapping> removedParameterMappings = new HashSet<>();
        List<UpdateSet> removedSet = new ArrayList<>();
        StringBuilder logMessage = new StringBuilder();
        MessageContainer messageContainer = new MessageContainer();
        // 多个update语句通过forEach标签拼接时，需要计算当前update语句的参数映射的起始位置
        calculateStartParameterMappingIndex(parameterMappings, index, multiParserRequest);
        for (UpdateSet updateSet : updateSets) {
            String columnName = updateSet.getColumns().get(0).getColumnName();
            Expression expression = updateSet.getExpressions().get(0);
            if (!(expression instanceof JdbcParameter)) {
                // 仅处理 预编译参数
                continue;
            }
            ParameterMapping parameterMapping = parameterMappings.get(multiParserRequest.getAndIncrement());
            String propertyName = parameterMapping.getProperty();
            Object property = null;
            try {
                // get property value
                property =
                        ParameterPropertyUtils.getPropertyValue(multiParserRequest, parameterMapping, parameterObject);
            } catch (Exception e) {
                log.warn(MessageFormat.format("Fail to get property: {0} in parameter: {1}, with exception: ",
                        propertyName, JsonUtils.toJsonString(parameterObject)), e);
                continue;
            }
            // match the desensitized date
            if (updateForbiddenMatcherComposite.match(propertyName, property, parameterObject, columnName,
                    messageContainer)) {
                logMessage.append(
                                MessageFormat.format("Removed update set, propertyName: {0}, property: {1}, columnName: {2}",
                                        propertyName, JsonUtils.toJsonString(property), columnName))
                        .append(CRLF);
                removedSet.add(updateSet);
                removedParameterMappings.add(parameterMapping);
            }
        }
        // if removedSet is not empty, then remove the set with desensitized data from updateSets
        if (CollectionUtils.isNotEmpty(removedSet)) {
            updateSets.removeAll(removedSet);
            multiParserRequest.setRemovedUpdate(true);
            // log and metric
            logMessage.append(MessageFormat.format("ParameterObject: {0}", JsonUtils.toJsonString(parameterObject)))
                .append(CRLF);

            Metrics.LOGGER.collect(new MetricMeta(SECURITY_DAL_FORBIDDEN + LOGGER_SUFFIX)
                .withTag(TABLE, ofNullable(update.getTable()).map(Table::getName).orElse(UNKNOWN))
                .withTag(REQUEST_URI, ofNullable(RequestContext.getCurrentContext().getRequestURI()).orElse(UNKNOWN))
                .withTag(FORBIDDEN_MATCHER,
                    ofNullable(messageContainer.getForbiddenMatcherNameSet()).map(TreeSet::toString).orElse(UNKNOWN))
                .setMessage(logMessage.toString()));
            Metrics.REGISTRY.timer(SHARD_FORBIDDEN_SINGLE
                .withTag(TABLE, ofNullable(update.getTable()).map(Table::getName).orElse(UNKNOWN))
                .withTag(REQUEST_URI, ofNullable(RequestContext.getCurrentContext().getRequestURI()).orElse(UNKNOWN))
                .withTag(FORBIDDEN_MATCHER,
                    ofNullable(messageContainer.getForbiddenMatcherNameSet()).map(TreeSet::toString).orElse(UNKNOWN)))
                .record(getDuration(currentStartTime));
        }
        if (CollectionUtils.isNotEmpty(removedParameterMappings)) {
            List<ParameterMapping> newParameterMappings = parameterMappings.stream()
                .filter(parameterMapping -> !removedParameterMappings.contains(parameterMapping))
                .collect(Collectors.toList());
            boundSql.parameterMappings(newParameterMappings);
            multiParserRequest.setParameterMappingIndex(
                multiParserRequest.getParameterMappingIndex() - removedParameterMappings.size());
        }
        if (CollectionUtils.isEmpty(updateSets) && CollectionUtils.isNotEmpty(removedSet)) {
            // if updateSets is empty after removed, then padding original update value to make sql legal
            UpdateSet updateSet = removedSet.get(0);
            Column column = CollectionUtils.isNotEmpty(updateSet.getColumns()) ? updateSet.getColumns().get(0)
                : new Column(update.getTable(), "");
            updateSets.add(new UpdateSet(column, column));
        }
    }

    /**
     * if insert sql contains desensitized data, then record related information in order to log and metric
     * 
     * @param insert the insert statement
     * @param index the index of multi sql
     * @param sql the sql
     * @param obj the parameter object
     */
    @Override
    protected void processInsert(Insert insert, int index, String sql, Object obj) {
        Instant currentStartTime = Instant.now();
        List<Column> columns = insert.getColumns();
        // validate the parameters
        if (CollectionUtils.isEmpty(columns)) {
            return;
        }
        if (!(obj instanceof MultiParserRequest)) {
            return;
        }
        MultiParserRequest multiParserRequest = (MultiParserRequest)obj;
        if (!IPluginPropertyParserRequest.validate(multiParserRequest)) {
            return;
        }
        DSBoundSql boundSql = multiParserRequest.getBoundSql();
        if (boundSql == null) {
            return;
        }

        Object parameterObject = boundSql.parameterObject();
        List<ParameterMapping> parameterMappings = boundSql.parameterMappings();
        if (parameterObject == null || CollectionUtils.isEmpty(parameterMappings)) {
            return;
        }
        MessageContainer messageContainer = new MessageContainer();
        Exception exception = null;
        try {
            // traverse
            ItemsList itemsList = insert.getItemsList();
            if (itemsList != null) {
                itemsList.accept(new ItemsListVisitor() {
                    @Override
                    public void visit(SubSelect subSelect) {
                        // empty method
                    }

                    @Override
                    public void visit(ExpressionList expressionList) {
                        processExpressionListForInsert(expressionList, parameterObject, parameterMappings, columns,
                            messageContainer, multiParserRequest);
                    }

                    @Override
                    public void visit(NamedExpressionList namedExpressionList) {
                        // empty method
                    }

                    @Override
                    public void visit(MultiExpressionList multiExpressionList) {
                        // 处理批量插入的SQL
                        List<ExpressionList> expressionLists = multiExpressionList.getExpressionLists();
                        if (CollectionUtils.isEmpty(expressionLists)) {
                            return;
                        }
                        expressionLists
                            .forEach(expressionList -> processExpressionListForInsert(expressionList, parameterObject,
                                parameterMappings, columns, messageContainer, multiParserRequest));
                    }
                });
            }
        } catch (Exception e) {
            exception = e;
        }
        // log and metric
        Metrics.REGISTRY
            .timer(SHARD_FORBIDDEN_INSERT
                .withTag(TABLE, ofNullable(insert.getTable()).map(Table::getName).orElse(UNKNOWN))
                .withTag(REQUEST_URI,
                    ofNullable(RequestContext.getCurrentContext().getRequestURI()).orElse(UNKNOWN))
                .withTag(FORBIDDEN_MATCHER,
                    ofNullable(messageContainer.getForbiddenMatcherNameSet()).map(TreeSet::toString)
                        .orElse(UNKNOWN))
                .withTag(SUCCESS_FLAG, exception == null))
            .record(getDuration(currentStartTime));
        if (exception != null) {
            messageContainer.getLogMessage().append(buildExceptionDetail(exception));
        }
        Metrics.LOGGER.collect(new MetricMeta(SHARD_FORBIDDEN_INSERT + LOGGER_SUFFIX)
            .withTag(REQUEST_URI, ofNullable(RequestContext.getCurrentContext().getRequestURI()).orElse(UNKNOWN))
            .withTag(FORBIDDEN_MATCHER,
                ofNullable(messageContainer.getForbiddenMatcherNameSet()).map(TreeSet::toString).orElse(UNKNOWN))
            .withTag(SUCCESS_FLAG, String.valueOf(exception == null))
            .setMessage(messageContainer.getLogMessage().toString()));
    }

    /**
     * get the String of exception stackTrace
     *
     * @param exception the exception
     */
    private static String buildExceptionDetail(Exception exception) {
        StringWriter sw = new StringWriter();
        PrintWriter pw = new PrintWriter(sw);
        exception.printStackTrace(pw);
        return MessageFormat.format(CRLF + "[{0}] Desensitization Process failed:" + CRLF + "{1}", "Insert Process",
            sw);
    }

    /**
     * process ExpressionList:
     * traverse the insert value，if match desensitization, then record related information in order to log and metric
     *
     * @param expressionList the List of expression
     * @param parameterObject the parameter Object
     * @param parameterMappings the parameterMappings
     * @param columns the columns
     * @param messageContainer the messageContainer which contains log message
     * @param multiParserRequest the extra request for parsing
     */
    private void processExpressionListForInsert(ExpressionList expressionList, Object parameterObject,
        List<ParameterMapping> parameterMappings, List<Column> columns, MessageContainer messageContainer,
        MultiParserRequest multiParserRequest) {
        List<Expression> expressions = expressionList.getExpressions();
        if (CollectionUtils.isEmpty(expressions)) {
            return;
        }
        if (expressions.size() != columns.size()) {
            log.warn("The itemsList size[{}] is not equal to columns size[{}]", expressions.size(),
                columns.size());
            return;
        }
        StringBuilder logMessage = messageContainer.getLogMessage();
        for (int i = 0; i < expressions.size(); i++) {
            Expression expression = expressions.get(i);
            if (!(expression instanceof JdbcParameter)) {
                // only process prepared parameter
                continue;
            }
            JdbcParameter jdbcParameter = (JdbcParameter)expression;
            Integer parameterIndex = jdbcParameter.getIndex();
            if (parameterIndex == null) {
                continue;
            }
            ParameterMapping parameterMapping = parameterMappings.get(parameterIndex - 1);
            String propertyName = parameterMapping.getProperty();
            Object property = null;
            try {
                property =
                    ParameterPropertyUtils.getPropertyValue(multiParserRequest, parameterMapping, parameterObject);
            } catch (Exception e) {
                log.warn(MessageFormat.format("Fail to get property: {0} in parameter: {1}, with exception: ",
                    propertyName, JsonUtils.toJsonString(parameterObject)), e);
                continue;
            }
            String columnName = Optional.ofNullable(columns.get(i)).map(Column::getColumnName).orElse(UNKNOWN);
            if (updateForbiddenMatcherComposite.match(propertyName, property, parameterObject,
                columnName,
                messageContainer)) {
                logMessage.append(
                    MessageFormat.format(
                        "Detected desensitized data in Insert SQL, propertyName: {0}, property: {1}, columnName: {2}",
                        propertyName, JsonUtils.toJsonString(property), columnName))
                    .append(CRLF);
            }
        }
    }

    /**
     * calculate the start index of ParameterMappings with multi sql
     * 
     * @param parameterMappings the parameterMappings
     * @param index the index of multi sql
     * @param multiParserRequest the extra request for parsing
     */
    private void calculateStartParameterMappingIndex(List<ParameterMapping> parameterMappings, int index,
        MultiParserRequest multiParserRequest) {
        // first sql no need calculate
        if (index == 0) {
            return;
        }
        Integer startIndex = multiParserRequest.getParameterMappingIndex();
        // skip the element which less than #index
        for (; startIndex < parameterMappings.size(); startIndex++) {
            ParameterMapping parameterMapping = parameterMappings.get(startIndex);
            String propertyName = parameterMapping.getProperty();
            Matcher listMatcher = EXTRACT_LIST_INFO_PATTERN.matcher(propertyName);
            if (!listMatcher.matches()) {
                return;
            }
            int number = Integer.parseInt(listMatcher.group(1));
            if (number >= index) {
                break;
            }
        }
        multiParserRequest.setParameterMappingIndex(startIndex);
    }

    /**
     * prepare the request for parsing
     * 
     * @param sh StatementHandler(可能是代理对象)
     * @param connection Connection
     * @param transactionTimeout transactionTimeout
     */
    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        Instant startTime = Instant.now();
        PluginUtils.DSStatementHandler dsSh = PluginUtils.dsStatementHandler(sh);
        MappedStatement ms = dsSh.mappedStatement();
        SqlCommandType sct = ms.getSqlCommandType();
        if (sct == SqlCommandType.UPDATE) {
            PluginUtils.DSBoundSql mpBs = dsSh.dsBoundSql();
            MultiParserRequest multiParserRequest = new MultiParserRequest(mpBs, dsSh.configuration());
            mpBs.sql(parserMulti(mpBs.sql(), multiParserRequest));
            log.debug("Update sql after UpdateForbiddenPlugin process: {}", mpBs.sql());

            Metrics.REGISTRY.timer(SHARD_FORBIDDEN_BATCH
                .withTag(REMOVED_FLAG, multiParserRequest.isRemovedUpdate()))
                .record(getDuration(startTime));
        } else if (sct == SqlCommandType.INSERT) {
            PluginUtils.DSBoundSql mpBs = dsSh.dsBoundSql();
            MultiParserRequest multiParserRequest = new MultiParserRequest(mpBs, dsSh.configuration());
            mpBs.sql(parserMulti(mpBs.sql(), multiParserRequest));
        }
    }

    @Override
    public void setProperties(Properties properties) {
        for (String stringPropertyName : properties.stringPropertyNames()) {
            String property = properties.getProperty(stringPropertyName);
            UpdateForbiddenMatcher updateForbiddenMatcher = ClassUtils.newInstance(property);
            updateForbiddenMatcherComposite.add(updateForbiddenMatcher);
        }
        updateForbiddenMatcherComposite.sort();
    }

    @Override
    public int order() {
        return -100;
    }

    @Data
    public static class MessageContainer {

        /**
         * the name of forbiddenMatcher which execute forbid
         */
        private TreeSet<String> forbiddenMatcherNameSet;
        /**
         * log message
         */
        private StringBuilder logMessage = new StringBuilder();
    }

    /**
     * request arguments for parserMulti method
     */
    @Getter
    @NoArgsConstructor
    @AllArgsConstructor
    public static class MultiParserRequest implements IPluginPropertyParserRequest {
        /**
         * boundSql for modifying
         */
        private PluginUtils.DSBoundSql boundSql;
        /**
         * mybatis Configuration for parsing
         */
        private Configuration configuration;

        /**
         * index of parameterMappings
         */
        private Integer parameterMappingIndex = 0;

        /**
         * whether removed set in update sql
         */
        private boolean removedUpdate = false;

        public MultiParserRequest(DSBoundSql dsBoundSql, Configuration configuration) {
            this.boundSql = dsBoundSql;
            this.configuration = configuration;
        }

        /**
         * get the index of parameterMappings and increment
         *
         * @return the index of parameterMappings
         */
        public Integer getAndIncrement() {
            return parameterMappingIndex++;
        }

        @Override
        public PluginUtils.DSBoundSql getBoundSql() {
            return boundSql;
        }

        @Override
        public Configuration getConfiguration() {
            return configuration;
        }

        public void setParameterMappingIndex(Integer parameterMappingIndex) {
            this.parameterMappingIndex = parameterMappingIndex;
        }

        public void setRemovedUpdate(boolean removedUpdate) {
            this.removedUpdate = removedUpdate;
        }

    }

    /**
     * get the duration since start time
     */
    private static Duration getDuration(Instant startTime) {
        return Duration.between(startTime, Instant.now());
    }
}
